# build_checklist.py
# usage: 从 tmp_newquestions.txt 构建检查清单
# python ./src/tools/build_checklist.py --des tmp_checklist.txt

import os
import sys
import json
import argparse
from typing import List, Dict, Any

# 添加项目根目录到路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)

from query.get_sql_answer import get_sql_answer

def parse_questions_file(file_path: str) -> List[Dict[str, Any]]:
    """
    解析 tmp_newquestions.txt 文件：
    - 顶格行为模板名
    - 缩进行（以空格或Tab开头）为JSON问题
    每个模板只取第一个问题。
    """
    entries = []
    current_template = None

    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, start=1):
            line = line.rstrip()  # 去除换行符
            if not line.strip():
                continue  # 跳过空行

            # 判断是否为模板行：无缩进（顶格）
            if not line.startswith(' ') and not line.startswith('\t'):
                current_template = line.strip()
                print(f"📌 模板行: {current_template}")
            else:
                # 缩进行：尝试解析为 JSON（问题）
                json_str = line.strip()
                if not json_str.startswith('{'):
                    continue  # 不是 JSON，跳过（如注释等）
                try:
                    question_data = json.loads(json_str)
                    # 只取每个模板的第一个问题
                    if current_template and not any(e['template'] == current_template for e in entries):
                        entries.append({
                            'template': current_template,
                            'question_data': question_data
                        })
                        print(f"✅ 提取问题: id={question_data['id']}, 问题='{question_data['question'][:40]}...'")
                except json.JSONDecodeError as e:
                    print(f"❌ 第 {line_num} 行 JSON 解析失败: {json_str} -> {e}")
                    continue

    return entries

def main():
    parser = argparse.ArgumentParser(description="从 tmp_newquestions.txt 构建检查清单")
    parser.add_argument('--des', required=True, help='输出结果文件路径')
    args = parser.parse_args()

    input_file = 'tmp_newquestions.txt'
    output_file = args.des

    # 解析问题
    questions = parse_questions_file(input_file)
    print(f"共找到 {len(questions)} 个不同的问题模板，开始处理...\n")

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_file), exist_ok=True) if os.path.dirname(output_file) else None

    # 逐条处理并实时写入
    with open(output_file, 'w', encoding='utf-8') as f_out:
        for idx, item in enumerate(questions, start=1):
            q_data = item['question_data']
            template = item['template']
            q_id = q_data['id']
            question = q_data['question']
            level = q_data['level']

            # 调用 get_sql_answer 获取 SQL
            try:
                sql = get_sql_answer(question)
                # 假设 get_sql_answer 返回字符串，若返回其他结构请调整
                if sql is None:
                    sql = ""
            except Exception as e:
                print(f"[ERROR] 调用 get_sql_answer 失败 (id={q_id}): {e}")
                sql = ""

            # 构造输出字典
            result = {
                "id": q_id,
                "flag": 0,
                "sql": sql,
                "template": template,
                "question": question,
                "level": level
            }

            # 实时输出到屏幕
            print(f"[{idx}/{len(questions)}] id={q_id}, question='{question[:50]}...'")
            # 立即写入文件（每条一行，JSON格式）
            f_out.write(json.dumps(result, ensure_ascii=False) + '\n')
            f_out.flush()  # 确保立即写入磁盘

    print(f"\n✅ 处理完成，结果已保存至: {output_file}")


if __name__ == '__main__':
    main()